#include    "socket.h"
#include    "logger.h"

#include	<stdexcept>

#if defined(_WIN32)
#	define	FD_SETSIZE	256
#	include	<WinSock2.h>
#	include	<WS2tcpip.h>
#	include	<map>
#	pragma	comment(lib, "ws2_32.lib")
#else
#	include	<cstdlib>
#	include	<stdexcept>
#	include	<arpa/inet.h>
#	include	<fcntl.h>
#	include	<netinet/in.h>
#	include	<sys/epoll.h>
#	include	<sys/socket.h>
#	include	<sys/types.h>
#	include	<unistd.h>
#	include	<cstring>
#endif

#if defined(_WIN32)
class IOModule {
public:
	IOModule() : _r_set{ 0 }, _w_set{ 0 }, _listeners() {}
	~IOModule() {}

	void Add(int fd, IOListener * listener, int events) {
		SOCKET socket = (SOCKET)fd;
		
		if (socket == INVALID_SOCKET) return;
		if ((events & IO_READ) && !FD_ISSET(socket, &_r_set)) FD_SET(socket, &_r_set);
		if ((events & IO_WRITE) && !FD_ISSET(socket, &_w_set)) FD_SET(socket, &_w_set);

		_listeners[fd] = listener;
	}

	void Wait(int fd, IOListener * listener, int events) {
		SOCKET socket = (SOCKET)fd;
		if (socket == INVALID_SOCKET) return;

		fd_set read_set, write_set;
		FD_ZERO(&read_set);
		FD_ZERO(&write_set);
		if (events & IO_READ) FD_SET(socket, &read_set);
		if (events & IO_WRITE) FD_SET(socket, &write_set);
		select(0, &read_set, &write_set, 0, NULL);

		if (FD_ISSET(socket, &read_set)) listener->OnReadable();
		if (FD_ISSET(socket, &write_set)) listener->OnWritable();
	}

	void Del(int fd) {
		SOCKET socket = (SOCKET)fd;
		if (FD_ISSET(socket, &_r_set)) FD_CLR(socket, &_r_set);
		if (FD_ISSET(socket, &_w_set)) FD_CLR(socket, &_w_set);
		auto it = _listeners.find(fd);
		if (it != _listeners.end()) _listeners.erase(it);
	}

	void Breath() {
		static fd_set cr_set = { 0 };
		static fd_set cw_set = { 0 };
		static struct timeval wait = { 0, 0 };

		int r_count = _r_set.fd_count;
		cr_set.fd_count = r_count;
		if (r_count > 0) memcpy(cr_set.fd_array, _r_set.fd_array, sizeof(SOCKET) * r_count);

		int w_count = _w_set.fd_count;
		cw_set.fd_count = w_count;
		if (w_count > 0) memcpy(cw_set.fd_array, _w_set.fd_array, sizeof(SOCKET) * w_count);

		if (select(0, &cr_set, &cw_set, NULL, &wait) <= 0) return;

		for (u_int n = 0; n < cr_set.fd_count; ++n) {
			auto it = _listeners.find((int)cr_set.fd_array[n]);
			if (it != _listeners.end()) it->second->OnReadable();
		}

		for (u_int n = 0; n < cw_set.fd_count; ++n) {
			auto it = _listeners.find((int)cw_set.fd_array[n]);
			if (it != _listeners.end()) it->second->OnWritable();
		}
	}

private:
	fd_set _r_set;
	fd_set _w_set;
	std::map<int, IOListener *> _listeners;
};
#else
class IOModule {
public:
	IOModule() : _epoll(epoll_create(1)) {}
	~IOModule() { if (_epoll >= 0) close(_epoll); }

	void Add(int fd, IOListener * listener, int events) {
		epoll_event ev;
		ev.data.ptr = listener;

		if (events & IO_READ) ev.events |= EPOLLIN;
		if (events & IO_WRITE) ev.events |= EPOLLOUT;

		epoll_ctl(_epoll, EPOLL_CTL_ADD, fd, &ev);
	}

	void Wait(int fd, IOListener * listener, int events) {
		int wait_fd = epoll_create(1);
		if (wait_fd < 0) return;

		epoll_event ev;
		ev.events = EPOLLET | EPOLLRDHUP;

		if (events & IO_READ) ev.events |= EPOLLIN;
		if (events & IO_WRITE) ev.events |= EPOLLOUT;

		epoll_ctl(wait_fd, EPOLL_CTL_ADD, fd, &ev);

		if (epoll_wait(wait_fd, &ev, 1, -1) > 0) {
			if (ev.events & EPOLLIN) listener->OnReadable();
			if (ev.events & EPOLLOUT) listener->OnWritable();
		}

		close(wait_fd);
	}

	void Del(int fd) {
		epoll_ctl(_epoll, EPOLL_CTL_DEL, fd, NULL);
	}

	void Breath() {
		static epoll_event events[256] = { 0 };

		int count = epoll_wait(_epoll, events, 256, 0);
		if (count <= 0) return;

		for (int i = 0; i < count; ++i) {
			IOListener * listener = (IOListener *)events[i].data.ptr;
			if (events[i].events & EPOLLIN) listener->OnReadable();
			if (events[i].events & EPOLLOUT) listener->OnWritable();
		}
	}

private:
	int _epoll;
};
#endif

static IOModule * GetIOModule() {
	static IOModule * g_io = new IOModule();
	return g_io;
}

Socket::Socket() {
#if defined(_WIN32)
	WSADATA data;
	if (WSAStartup(MAKEWORD(2, 2), &data)) throw std::runtime_error("WSAStartup failed!!!");
	ctx = (int)INVALID_SOCKET;
#else
	ctx = -1;
#endif
}

Socket::~Socket() {
	Close();
#if defined(_WIN32)
	WSACleanup();
#endif
}

bool Socket::Connect(const std::string & ip, int port) {
#if defined(_WIN32)
	ASSERT_RET(ctx == (int)INVALID_SOCKET, false, "Socket::Connect failed. Already connected!");
#else
	ASSERT_RET(ctx < 0, false, "Socket::Connect failed. Already connected!");
#endif

	struct sockaddr_in addr;
	memset(&addr, 0, sizeof(addr));

	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);

	ASSERT_RET(inet_pton(AF_INET, ip.c_str(), &addr.sin_addr) == 1, false, "Socket::Connect failed. Bad remote address : %s!", ip.c_str());

#if defined(_WIN32)
	SOCKET fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	ASSERT_RET(fd != INVALID_SOCKET, false, "Socket::Connect failed due to create socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)");

	u_long flags = 1;
	if (ioctlsocket(fd, FIONBIO, &flags) != 0) {
		closesocket(fd);
		LOG_ERR("Socket::Connect failed due to ioctlsocket(fd, FIONBIO, 1)");
		return false;
	}
#else
	int fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	ASSERT_RET(fd >= 0, false, "Socket::Connect failed due to create socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)");

	int flags = fcntl(fd, F_GETFL, 0);
	if (fcntl(fd, F_SETFL, O_NONBLOCK | flags) < 0) {
		close(fd);
		LOG_ERR("Socket::Connect failed due to fcntl(fd, F_SETFL, O_NONBLOCK | flags)");
		return false;
	}
#endif

	if (connect(fd, (sockaddr *)&addr, sizeof(sockaddr)) < 0) {
#if defined(_WIN32)
		int err = WSAGetLastError();
		if (err != WSAEWOULDBLOCK) {
			closesocket(fd);
			LOG_ERR("Socket::Connect failed with errcode : %d", err);
			return false;
		}

		fd_set wait_set;
		struct timeval wait_time;

		wait_time.tv_sec = 1;
		wait_time.tv_usec = 0;

		FD_ZERO(&wait_set);
		FD_SET(fd, &wait_set);

		if (select(0, 0, &wait_set, 0, &wait_time) <= 0) {
			closesocket(fd);
			LOG_ERR("Socket::Connect failed due to timeout");
			return false;
		}
#else
		if (errno != EINPROGRESS) {
			close(fd);
			LOG_ERR("Socket::Connect failed with errcode : %d", errno);
			return false;
		}

		int wait_fd = epoll_create(1);
		ASSERT_RET(wait_fd >= 0, false, "Socket::Connect failed due to epoll_create(1) for async waiting");

		epoll_event ev;
		ev.events = EPOLLET | EPOLLOUT | EPOLLRDHUP;
		epoll_ctl(wait_fd, EPOLL_CTL_ADD, fd, &ev);

		int count = epoll_wait(wait_fd, &ev, 1, 1000);
		close(wait_fd);
		if (count < 0 || ev.events & EPOLLRDHUP) {
			close(fd);
			LOG_ERR("Socket::Connect failed due to timeout");
			return false;
		}
#endif
	}

	ctx = (int)fd;
	return true;
}

bool Socket::Listen(const std::string & ip, int port) {
#if defined(_WIN32)
	ASSERT_RET(ctx == (int)INVALID_SOCKET, false, "Socket::Listen failed. Already connected!");
#else
	ASSERT_RET(ctx < 0, false, "Socket::Listen failed. Already connected!");
#endif

	struct sockaddr_in addr;
	memset(&addr, 0, sizeof(addr));

	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);

	ASSERT_RET(inet_pton(AF_INET, ip.c_str(), &addr.sin_addr) == 1, false, "Socket::Listen failed. Bad listen address : %s!", ip.c_str());

#if defined(_WIN32)
	SOCKET fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	ASSERT_RET(fd != INVALID_SOCKET, false, "Socket::Listen failed due to create socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)");

	u_long flags = 1;
	if (ioctlsocket(fd, FIONBIO, &flags) != 0) {
		closesocket(fd);
		LOG_ERR("Socket::Listen failed due to ioctlsocket(fd, FIONBIO, 1)");
		return false;
	}

	if (::bind(fd, (sockaddr *)&addr, sizeof(sockaddr)) < 0 || ::listen(fd, FD_SETSIZE) < 0) {
		closesocket(fd);
		LOG_ERR("Socket::Listen failed due to bind() or listen()");
		return false;
	}
#else
	int fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	ASSERT_RET(fd >= 0, false, "Socket::Listen failed due to create socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)");

	int flags = fcntl(fd, F_GETFL, 0);
	if (fcntl(fd, F_SETFL, O_NONBLOCK | flags) < 0) {
		close(fd);
		LOG_ERR("Socket::Listen failed due to fcntl(fd, F_SETFL, O_NONBLOCK | flags)");
		return false;
	}

	int reuse = 1;
	setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse));

	if (::bind(fd, (sockaddr *)&addr, sizeof(sockaddr)) < 0 || ::listen(fd, 1024) < 0) {
		close(fd);
		LOG_ERR("Socket::Listen failed due to bind() or listen()");
		return false;
	}
#endif	

	ctx = (int)fd;
	return true;
}

Socket * Socket::Accept() {
#if defined(_WIN32)
	if (ctx == (int)INVALID_SOCKET) return nullptr;

	sockaddr_in addr = { 0 };
	int addr_size = sizeof(addr);
	SOCKET ctx_accept = accept((SOCKET)ctx, (sockaddr *)&addr, &addr_size);
	if (ctx_accept == INVALID_SOCKET) return nullptr;

	u_long flags = 1;
	if (ioctlsocket(ctx_accept, FIONBIO, &flags) != 0) {
		closesocket(ctx_accept);
		return nullptr;
	}

	Socket * socket = new Socket;
	socket->ctx = (int)ctx_accept;
	return socket;
#else
	if (ctx < 0) return nullptr;

	sockaddr_in addr = { 0 };
	socklen_t addr_size = sizeof(addr);
	int fd = accept(ctx, (sockaddr *)&addr, &addr_size);
	if (fd < 0) return nullptr;

	int flags = fcntl(fd, F_GETFL, 0);
	if (fcntl(fd, F_SETFL, O_NONBLOCK | flags) < 0) {
		close(fd);
		return nullptr;
	}

	Socket * socket = new Socket;
	socket->ctx = fd;
	return socket;
#endif	
}

int Socket::Recv(char * buf, size_t size) {
#if defined(_WIN32)
	if (ctx == (int)INVALID_SOCKET) return -1;
	int recved = recv((SOCKET)ctx, buf, (int)size, 0);
	if (recved > 0) return recved;
	if (recved < 0 && WSAGetLastError() == WSAEWOULDBLOCK) return 0;
	return recved == 0 ? -1 : -2;
#else
	if (ctx < 0) return -1;
	int recved = (int)recv(ctx, buf, (int)size, MSG_DONTWAIT);
	if (recved > 0) return recved;
	if (recved < 0 && errno == EAGAIN) return 0;
	return recved == 0 ? -1 : -2;
#endif	
}

bool Socket::Send(const char * data, size_t size) {
#if defined(_WIN32)
	if (ctx == (int)INVALID_SOCKET) return false;

	char * ptr = (char *)data;
	int sent = 0;
	int left = (int)size;

	while (true) {
		sent = send((SOCKET)ctx, ptr, left, 0);
		if (sent < 0) {
			if (WSAGetLastError() == WSAEWOULDBLOCK) {
				Sleep(1);
			} else {
				return false;
			}
		} else if (sent < left) {
			left -= sent;
			ptr += sent;
		} else if (sent == left) {
			return true;
		} else {
			return left == 0;
		}
	}
#else
	if (ctx < 0) return false;

	char * ptr = (char *)data;
	int sent = 0;
	int left = (int)size;

	while (true) {
		sent = (int)send(ctx, ptr, left, MSG_DONTWAIT);
		if (sent < 0) {
			if (errno == EAGAIN) {
				usleep(1000);
			} else {
				return false;
			}
		} else if (sent < left) {
			left -= sent;
			ptr += sent;
		} else if (sent == left) {
			return true;
		} else {
			return left == 0;
		}
	}
#endif
}

void Socket::Close() {
#if defined(_WIN32)
	if (ctx == (int)INVALID_SOCKET) return;
	closesocket((SOCKET)ctx);
	ctx = (int)INVALID_SOCKET;
#else
	if (ctx < 0) return;
	close(ctx);
	ctx = -1;
#endif
}

void IOListener::BindIO(int fd, int events) {
	GetIOModule()->Add(fd, this, events);
	_binded_io_fd = fd;
}

void IOListener::WaitIO(int event) {
	GetIOModule()->Wait(_binded_io_fd, this, event);
}

void IOListener::UnbindIO() {
	GetIOModule()->Del(_binded_io_fd);
}

void BreathAllSockets() {
	GetIOModule()->Breath();
}